import os
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset, random_split
from transformers import AutoTokenizer
from datasets import load_dataset,Dataset
"""
包含了从文件，目录，集合，自定义Dataset等形式来加载数据
"""
def read_data_by_csv():
    df = pd.read_csv("data/ChnSentiCorp_htl_all.csv")
    df = df.dropna()
    return df


class MyDataset(Dataset):
    def __init__(self)->None:
        super().__init__()
        # self.data = pd.read_csv("data/ChnSentiCorp_htl_all.csv")
        # self.data = self.data.dropna()
        self.data = read_data_by_csv()

    def __getitem__(self, index):
        return self.data.iloc[index]["review"], self.data.iloc[index]["label"]

    def __len__(self):
        return len(self.data)

def get_dataset_by_dir():
    """"""
    dataset = load_dataset("csv", data_dir="data/", split="train")
    return dataset

def get_dataset_by_file():
    dataset = load_dataset("csv", data_files="data/ChnSentiCorp_htl_all.csv", split="train")
    return dataset

def get_dataset_by_pandas():
    df = pd.read_csv("data/ChnSentiCorp_htl_all.csv")
    df = df.dropna()
    dataset = Dataset.from_pandas(df)
    return dataset

def get_dataset_by_list():
    data_dict = [
        {
            'label': 1,
            'review': "距离川沙公路较近,但是公交指示不对,如果是\"蔡陆线\"的话,会非常麻烦.建议用别的路线.房间较为简单."
        },
        {
            'label': 1,
            'review': "商务大床房，房间很大，床有2M宽，整体感觉经济实惠不错!"
        },
        {
            'label': 1,
            'review': "早餐太差，无论去多少人，那边也不加食品的。酒店应该重视一下这个问题了。房间本身很好。"
        },
        {
            'label': 1,
            'review': "CBD中心,周围没什么店铺,说5星有点勉强.不知道为什么卫生间没有电吹风"
        },
        {
            'label': 1,
            'review': "总的来说，这样的酒店配这样的价格还算可以，希望他赶快装修，给我的客人留些好的印象"
        },
        {
            'label': 1,
            'review': "价格比比较不错的酒店。这次免费升级了，感谢前台服务员。房子还好，地毯是新的，比上次的好些。早餐的人很多要早去些。"
        },
        {
            'label': 1,
            'review': "不错，在同等档次酒店中应该是值得推荐的！"
        },
        {
            'label': 1,
            'review': "早餐很丰富，服务也热情，早上很早退房时，前台值此人员办理手续也非常快．"
        },
        {
            'label': 1,
            'review': "价格偏高,服务一般.窗外风景ok.早餐还不错"
        },
        {
            'label': 1,
            'review': "沈阳市政府的酒店，比较大气，交通便利，出门往左就是北陵公园，环境好。"
        }
    ]
    dataset = Dataset.from_list(data_dict)
    return dataset

def get_dataset_by_dict():
    data_dict = {
        'label': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        'review': [
            "商务大床房，房间很大，床有2M宽，整体感觉经济实惠不错!",
            "早餐太差，无论去多少人，那边也不加食品的。酒店应该重视一下这个问题了。房间本身很好。",
            "CBD中心,周围没什么店铺,说5星有点勉强.不知道为什么卫生间没有电吹风",
            "总的来说，这样的酒店配这样的价格还算可以，希望他赶快装修，给我的客人留些好的印象",
            "价格比比较不错的酒店。这次免费升级了，感谢前台服务员。房子还好，地毯是新的，比上次的好些。早餐的人很多要早去些。",
            "不错，在同等档次酒店中应该是值得推荐的！",
            "价格偏高,服务一般.窗外风景ok.早餐还不错",
            "沈阳市政府的酒店，比较大气，交通便利，出门往左就是北陵公园，环境好。",
            "早餐很丰富，服务也热情，早上很早退房时，前台值此人员办理手续也非常快．",
            "早餐很丰富，服务也热情，早上很早退房时，前台值此人员办理手续也非常快．"
        ]
    }
    dataset = Dataset.from_dict(data_dict)
    return dataset

def get_splited_dataset():
    dataset = get_dataset_by_file()
    print(f"num_rows of dataset:{dataset.num_rows}")
    trainset, validset = random_split(dataset=dataset,lengths=[0.9, 0.1])
    return trainset, validset

def test_dataset():
    dataset = MyDataset()
    for i in range(5):
        print(dataset[i])
    print("=================================")
    trainset, validset = get_splited_dataset()
    print(f"len of trainset:{len(trainset)}")
    print(f"len of validset:{len(validset)}")
    for i in range(10):
        print(trainset[i])

def get_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained('uer/roberta-base-finetuned-dianping-chinese')
    return tokenizer

def collate_fn(batch):
    tokenizer = get_tokenizer()
    texts, labels = [], []
    for item in batch:
        print(f"===item:{item}")
        """如果使用自定义的MyDataset使用下面数组item[0]、item[1]"""
        # texts.append(item[0])
        # labels.append(item[1])
        """如果使用load_dataset加载的数据就使用下面这种形式"""
        texts.append(item["review"])
        labels.append(item["label"])

    print(f"====texts:{texts}")
    inputs = tokenizer(texts,
                       max_length=128,
                       padding="max_length",
                       truncation=True,
                       return_tensors="pt")
    inputs["labels"] = torch.tensor(labels)
    return inputs

def get_dataloader():
    trainset, validset = get_splited_dataset()
    print(f"trainset:{trainset}")
    print(f"validset:{validset}")

    train_dataloader = DataLoader(trainset, batch_size=32, shuffle=True, collate_fn=collate_fn)
    valid_dataloader = DataLoader(validset, batch_size=64, shuffle=False, collate_fn=collate_fn)
    return train_dataloader, valid_dataloader

def test_dataloader():
    train_dataloader, valid_dataloader = get_dataloader()
    re = next(enumerate(valid_dataloader))[1]
    print("-----------------------------------")
    print(re)
    print("-----------------------------------")

if __name__ == "__main__":
    os.environ['HTTP_PROXY'] = 'http://127.0.0.1:10792'
    os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:10792'

    # test_dataset()
    test_dataloader()